{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dDtUTck7l0Kc",
        "outputId": "7632514b-84ab-4840-b3dc-97e953958ed9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mounted at /content/gdrive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "b-el3MsdoOX_"
      },
      "outputs": [],
      "source": [
        "!pip install ultralytics\n",
        "from IPython import display\n",
        "display.clear_output()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "WvgxJTMpokZ_"
      },
      "outputs": [],
      "source": [
        "from ultralytics import YOLO"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "vGJuFtTMoq-l"
      },
      "outputs": [],
      "source": [
        "model = YOLO(\"/content/gdrive/MyDrive/Weapons_Detection_Model/WD_Model/weights/best.pt\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YTEUOtpbtP56"
      },
      "outputs": [],
      "source": [
        "# Image Test\n",
        "result = model.predict(source=\"/content/test.jpg\", save=True)  # to specify certain confidence on prediction add conf argument (e.g. conf=0.4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mWl45T9NrFpx"
      },
      "outputs": [],
      "source": [
        "# Video Test\n",
        "result = model.predict(source=\"/content/test.mp4\", save=True) # (conf=0.4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5TabFwhMjKnP"
      },
      "outputs": [],
      "source": [
        "# Live detection test (only works locally)\n",
        "result = model.predict(source=0, save=True) # (conf=0.4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Filtered Images (Evidences)\n",
        "detections = [res for res in result if res.boxes]\n",
        "\n",
        "os.makedirs('/content/Evidences', exist_ok=True)\n",
        "\n",
        "if detections:\n",
        "    for det in detections:\n",
        "        original_filename = os.path.basename(det.path)   \n",
        "        det.save(f'/content/Evidences/{original_filename}') \n",
        "    print(\"Media files with detections saved to Evidences folder.\")\n",
        "else:  \n",
        "    print(\"No detections found.\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
